import os
import random

def read_and_split_by_label(file_path):
    pos_samples = []
    neg_samples = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            if line.strip():
                label = line.strip().split("\t")[0]
                if label == '1':
                    pos_samples.append(line)
                elif label == '0':
                    neg_samples.append(line)
    return pos_samples, neg_samples

def shrink_and_save(file_path, output_path, target_size_mb=10):
    target_bytes = target_size_mb * 1024 * 1024

    pos_samples, neg_samples = read_and_split_by_label(file_path)
    random.shuffle(pos_samples)
    random.shuffle(neg_samples)

    # 确保正负数量一致
    min_len = min(len(pos_samples), len(neg_samples))
    pos_samples = pos_samples[:min_len]
    neg_samples = neg_samples[:min_len]

    combined = []
    total_bytes = 0
    for p, n in zip(pos_samples, neg_samples):
        for line in (p, n):
            line_bytes = len(line.encode('utf-8'))
            if total_bytes + line_bytes > target_bytes:
                break
            combined.append(line)
            total_bytes += line_bytes

    # 写入缩小后的文件
    with open(output_path, 'w', encoding='utf-8') as f:
        f.writelines(combined)

    print(f"{output_path} 写入完成，总大小约为 {total_bytes / (1024 * 1024):.2f} MB，样本数：{len(combined)}")

if __name__ == "__main__":
    shrink_and_save("train.txt", "train.txt", target_size_mb=10)
    shrink_and_save("test.txt", "test.txt", target_size_mb=10)
